{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5377abc3-3633-4af3-89b0-095b450e55ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "\n",
    "import numpy as np\n",
    "import scqubits as scq\n",
    "\n",
    "\n",
    "from qiskit_metal.analyses.quantization.lumped_capacitive import load_q3d_capacitance_matrix\n",
    "from qiskit_metal.analyses.quantization.lom_core_analysis import CompositeSystem, Cell, Subsystem\n",
    "                        \n",
    "from scipy.constants import speed_of_light as c_light\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5723516-ee06-4a8a-bb99-f47a163f5df1",
   "metadata": {},
   "source": [
    "### 1. load fluxonium cell Q3d simulation results\n",
    "\n",
    "Loading the Maxwell capacitance matrices for the design as shown in the screenshot below:\n",
    "\n",
    "<img src=\"fluxonium_transmon_coupled.png\" width=600 height=600 />\n",
    "\n",
    "\n",
    "where we have a transmon coupled to a fluxonium through a direct coupler. \n",
    "\n",
    "For a simple introduction on Maxwell capacitance matrix, check out the following resources:\n",
    "https://www.fastfieldsolvers.com/Papers/The_Maxwell_Capacitance_Matrix_WP110301_R02.pdf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e02eb89-8f17-423d-b950-a0ba41073c29",
   "metadata": {},
   "outputs": [],
   "source": [
    "path = './Fluxonium_8p5MHz_cmat.txt'\n",
    "flux_mat, _, _, _ = load_q3d_capacitance_matrix(path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "485fe43f-be2d-4241-b146-e950b60c41c7",
   "metadata": {},
   "source": [
    "### load transmon cell Q3d simulation results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb74dda4-9751-4540-888f-e9b78a91b0f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "path = './Transmon_5p5GHz_fQ_cmat.txt'\n",
    "transmon_mat, _, _, _ = load_q3d_capacitance_matrix(path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7aff060f-2563-45b2-8bbc-d573bd663355",
   "metadata": {},
   "source": [
    "### 2. Create LOM cells from capacitance matrices\n",
    "#### Setting cell objects corresponding to the capacitance simulation results \n",
    "`coupler_pad_Q1` and `coupler_pad_Q2` refer to the same node corresponding to the direct coupler between the qubits but are different names in the capacitance matrix results file. In order to merge the two capacitance matrices in the LOM analysis, we need to rename them to be the same name. \n",
    "\n",
    "The following three parameters, `ind_dict`, `jj_dict`, `cj_dict`, all have the same structure. Each is a dictionary where the keys are tuples, giving the nodes that a junction is in between, and the values specifying the relevant values associated with the junction. `ind_dict` lets you specify the junction inductance in nH; `jj_dict` specifies the Josephson junction name (you can give the junction any name you wish; just need to be consistent with the name); `cj_dict` specifies the junction capacitance in fF. In the case of the fluxonium, we will set $E_j$ and $E_l$ directly later instead of deriving from the junction inductance; and since we are only concerned with capacitive coupling here (what's currently supported), `ind_dict` can just be a placeholder whose actual value is not important."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bfdd14e8-c4f8-42b2-81b6-b62a0bb9aed3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# cell 1\n",
    "\n",
    "opt1 = dict(\n",
    "    node_rename = {'coupling_pad_Q1': 'coupling'}, \n",
    "    cap_mat = flux_mat,\n",
    "    ind_dict = {('pad_bot_Q1', 'pad_top_Q1'): 1},  # placeholder inductance here; only used for node-basis transformation and reduction\n",
    "    jj_dict = {('pad_bot_Q1', 'pad_top_Q1'):'j1'},\n",
    ")\n",
    "cell_1 = Cell(opt1)\n",
    "\n",
    "# cell 2\n",
    "opt2 = dict(\n",
    "    node_rename = {'coupling_pad_Q2': 'coupling'},\n",
    "    cap_mat = transmon_mat,\n",
    "    ind_dict = {('pad_bot_Q2', 'pad_top_Q2'): 12.31},\n",
    "    jj_dict = {('pad_bot_Q2', 'pad_top_Q2'):'j2'},\n",
    ")\n",
    "cell_2 = Cell(opt2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "025bedb7-0a90-491d-a959-4f353abf7f86",
   "metadata": {},
   "source": [
    "### 3. Create subsystems\n",
    "#### Creating the four subsystems, corresponding to the 2 qubits\n",
    "\n",
    "`Subsystem` takes three required arguments. The four currently supported system types are `TRANSMON`, `FLUXONIUM`, `TL_RESONATOR` (transmission line resonator) and `LUMPED_RESONATOR`. `nodes` lets you specify which node the subsystem should be mapped to in the cells. They should be consistent with the node names you have given previously. `q_opts` lets specify any optional parameters you want to give. For example, for the fluxonium, you can provide `scqubits` parameters such as `EJ`, `EL` and `flux` here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4232478-69a1-49f0-882f-eedc4cc419ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "# subsystem 1: fluxonium\n",
    "fluxonium = Subsystem(name='fluxonium', sys_type='FLUXONIUM', nodes=['j1'], q_opts={'EJ':4860, 'EL':1140, 'flux': .5})\n",
    "\n",
    "\n",
    "# subsystem 2: transmon\n",
    "transmon = Subsystem(name='transmon', sys_type='TRANSMON', nodes=['j2'], q_opts={'ncut': 150, 'truncated_dim':10})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c7a0b33f-9be5-4ec4-a95c-1d5cb0ce0c03",
   "metadata": {},
   "source": [
    "### 4. Creat the composite system from the cells and the subsystems"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "141295eb-454c-47d1-b71d-08df911cfa2b",
   "metadata": {},
   "outputs": [],
   "source": [
    "composite_sys = CompositeSystem(\n",
    "    subsystems=[fluxonium, transmon], \n",
    "    cells=[cell_1, cell_2], \n",
    "    grd_node='ground_main_plane')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6797ebb5-aad6-45d9-b25c-db2d03926db6",
   "metadata": {},
   "source": [
    "The `circuitGraph` object encapsulates the lumped model circuit analysis (i.e., LOM analysis) and contain the intermediate as well as final L and C matrices, their inverses needed to construct the Hamiltonian of the composite system. For more details on the meaning and calculation of these matrices, check out https://arxiv.org/pdf/2103.10344.pdf.\n",
    "\n",
    "Just to note that you can use the analysis without needing to know any detail about this object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e32032b-31e4-4884-a695-9c89c629e689",
   "metadata": {},
   "outputs": [],
   "source": [
    "cg = composite_sys.circuitGraph()\n",
    "print(cg)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e85c5ca-d0e5-4a71-9f21-3d37ea5c82dd",
   "metadata": {},
   "source": [
    "### 5. Generate the hilberspace from the composite system, leveraging the scqubits package\n",
    "\n",
    "`add_interaction()` adds the interaction terms between the subsystems. Currently, capacitive coupling is supported (which is extracted by from off-diagonal elements in the C matrices, see *eqn 12, 13* in https://arxiv.org/pdf/2103.10344.pdf ) and contribute to the interaction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95171a27-702c-4e56-a13f-a8eca3e9ddba",
   "metadata": {},
   "outputs": [],
   "source": [
    "hilbertspace = composite_sys.create_hilbertspace()\n",
    "hilbertspace = composite_sys.add_interaction()\n",
    "print(hilbertspace)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "830e00b7-5908-4802-b177-432c402902e0",
   "metadata": {},
   "source": [
    "### 6. Print the results\n",
    "\n",
    "Print the calculated Hamiltonian parameters from diagonalized composite system Hamiltonian.\n",
    "\n",
    "The diagonal elements of the $\\chi$ matrix are the anharmonicities of the respective subsystems and the off-diagonal the dispersive shifts between them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2c91444-e60a-46e0-8fa6-83fb26b33f82",
   "metadata": {},
   "outputs": [],
   "source": [
    "hamiltonian_results = composite_sys.hamiltonian_results(hilbertspace, evals_count=30)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01403db9-0fa6-4f9b-b44e-602fb2804dca",
   "metadata": {},
   "outputs": [],
   "source": [
    "hamiltonian_results['chi_in_MHz'].to_dataframe()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b436a025-afb6-451c-891b-815deeea314c",
   "metadata": {},
   "source": [
    "The $\\chi$'s between the subsystems are based on the coupling strengths, $\\it{g}$'s between them (which are computed using the coupling capacitance (currently capacitive coupling is supported) and zero point fluctuations of the subsystem's charge operator at the coupling location)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9afefb68-9635-4762-8fe6-67b03fe8c8f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "composite_sys.compute_gs()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a9306ff-55f3-4ce2-9e63-4b348f2732d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "fluxonium.h_params"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8c53709-345b-453c-a105-8a3485873204",
   "metadata": {},
   "outputs": [],
   "source": [
    "transmon.h_params"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a3730af-bbc0-44cb-98e8-6a6c65e4ca4d",
   "metadata": {},
   "source": [
    "### 7. let's sweep some parameters now\n",
    "\n",
    "Let's sweep the flux from 0 to 1 in a grid of 100 points in unit of flux quantum using `scQubits`'s sweeping library."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ae6091a-b219-4543-bc99-d095d5ca3fc9",
   "metadata": {},
   "outputs": [],
   "source": [
    " _sys = hilbertspace.subsys_list[0] # fluxonium\n",
    "\n",
    "def update_hilbertspace(param_val):\n",
    "    _sys.flux = param_val\n",
    "    \n",
    "param_name = 'flux'\n",
    "param_vals = np.linspace(0, 1, 101)\n",
    "\n",
    "sweep = scq.ParameterSweep(\n",
    "    paramvals_by_name={param_name: param_vals},\n",
    "    evals_count=30,\n",
    "    hilbertspace=hilbertspace,\n",
    "    subsys_update_info={param_name: [_sys]},\n",
    "    update_hilbertspace=update_hilbertspace,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d295afd7-2ad8-495f-a97a-d8406c86cd9b",
   "metadata": {},
   "source": [
    "#### Plot transition frequencies as a function of the flux\n",
    "\n",
    "0->1 transition for the transmon\n",
    "\n",
    "0->1 transition for the fluxonium\n",
    "\n",
    "0->2 transition for the fluxonium"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ae5b901-987e-44b9-a162-389391c849f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "wq_t = sweep.transitions(False, [],(0, 0), (0, 1))[1][0]\n",
    "wq_f = sweep.transitions(False, [],(0, 0), (1, 0))[1][0]\n",
    "wq_f_02 = sweep.transitions(False, [],(0, 0), (2, 0))[1][0]\n",
    "\n",
    "\n",
    "plt.figure(figsize=(15, 6))\n",
    "plt.plot(param_vals, wq_t, 'ob', label='transmon')\n",
    "plt.plot(param_vals, wq_f, 'g-', label='fluxonium')\n",
    "plt.plot(param_vals, wq_f_02, 'r-', label='fluxonium_02')\n",
    "\n",
    "\n",
    "plt.xticks(param_vals[::10], rotation=45)\n",
    "\n",
    "plt.xlabel(r'flux ($\\Phi_{ext}/\\Phi_0$)')\n",
    "plt.ylabel(r'qubit freq (MHz)')\n",
    "plt.ylim([5550, 5650])\n",
    "\n",
    "plt.grid()\n",
    "plt.legend(fontsize=14)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ae71c91-f066-4060-ad22-e7ddd6154012",
   "metadata": {},
   "source": [
    "#### The dispersive shift, $\\chi$ between the two qubits as a function of the flux"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4e7c261-250f-46bb-a835-adbd0e9d7a0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "wq_f = sweep.transitions(False, [],(0, 0), (1, 0))[1][0]\n",
    "wq_f_t = sweep.transitions(False, [],(0, 1), (1, 1))[1][0]\n",
    "\n",
    "\n",
    "chi = wq_f_t - wq_f\n",
    "\n",
    "plt.figure(figsize=(15, 6))\n",
    "plt.plot(param_vals, chi, 'ob')\n",
    "\n",
    "plt.xticks(param_vals[::10], rotation=45)\n",
    "\n",
    "plt.xlabel(r'flux ($\\Phi_{ext}/\\Phi_0$)')\n",
    "plt.ylabel(r'$\\chi$ (MHz)')\n",
    "#plt.ylim([5550, 5650])\n",
    "\n",
    "plt.grid()\n",
    "plt.legend(fontsize=14)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "037d0350-1992-47b2-8695-fece3c60ced2",
   "metadata": {},
   "source": [
    "#### Zooming on the fluxonium sweet spot: its 0->1 transition as a function of the flux"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3715cbde-796c-4ed7-99b2-ca609d6ad575",
   "metadata": {},
   "outputs": [],
   "source": [
    "wq_f = sweep.transitions(False, [],(0, 0), (1, 0))[1][0]\n",
    "\n",
    "\n",
    "plt.figure(figsize=(15, 6))\n",
    "plt.plot(param_vals, wq_f, 'ob')\n",
    "\n",
    "\n",
    "plt.xticks(param_vals[::10], rotation=45)\n",
    "\n",
    "plt.xlabel(r'flux ($\\Phi_{ext}/\\Phi_0$)')\n",
    "plt.ylabel(r'fluxonium freq (MHz)')\n",
    "\n",
    "plt.grid()\n",
    "plt.legend(fontsize=14)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8037faeb-53f4-440c-abd7-df328884175a",
   "metadata": {},
   "outputs": [],
   "source": [
    "scq.get_units()\n",
    "scq.set_units('MHz')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34cb07fd-89ea-4117-9183-8f6c7251a910",
   "metadata": {},
   "source": [
    "#### Coherences of the fluxonium as a function of the flux\n",
    "\n",
    "For more information on their calculations and assumptions made, check out `scqubits` documentations:\n",
    "https://scqubits.readthedocs.io/en/latest/guide/guide-noise.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14457ed7-9be9-4746-b411-0b4a74094a16",
   "metadata": {},
   "outputs": [],
   "source": [
    "_sys.plot_t1_effective_vs_paramvals(param_name='flux', \n",
    "                                    param_vals=param_vals)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37e3a1ab-cfe8-4f86-9f66-1e74ed74b4d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "_sys.plot_coherence_vs_paramvals(param_name='flux', param_vals=param_vals)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b724859-ee14-449a-b815-013a700c34b4",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "metal_kernel",
   "language": "python",
   "name": "metal_kernel"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
